"""知识图谱工具模块"""
from typing import Dict, Any, List, Union
from langchain.tools import BaseTool
from langchain_core.tools import tool
from src.containers import get_knowledge_graph_service
import os

class KnowledgeGraphQueryTool(BaseTool):
    """知识图谱查询工具"""
    
    name: str = "knowledge_graph_query"
    description: str = """
    查询知识图谱获取结构化知识。适用于：
    - 查找实体之间的关系
    - 获取概念的定义和关联信息
    - 支持复杂推理查询
    - 增强现有搜索能力
    """
    
    def _run(self, query: str) -> str:
        """执行知识图谱查询"""
        kg_service = get_knowledge_graph_service()
        results = kg_service.hybrid_search(query)
        
        if not results:
            return "知识图谱中没有找到相关信息。"
        
        # 格式化结果
        formatted_results = []
        for result in results:
            formatted_result = f"{result['source']} --{result['relationship']}--> {result['target']}"
            if result.get('source_description'):
                formatted_result += f"\n  源描述: {result['source_description']}"
            if result.get('target_description'):
                formatted_result += f"\n  目标描述: {result['target_description']}"
            formatted_results.append(formatted_result)
        
        return "\n\n".join(formatted_results)
    
    async def _arun(self, query: str) -> str:
        """异步执行知识图谱查询"""
        return self._run(query)

@tool
def query_knowledge_graph(query: str) -> str:
    """查询知识图谱获取结构化知识和关系信息"""
    kg_service = get_knowledge_graph_service()
    results = kg_service.hybrid_search(query)
    
    if not results:
        return "知识图谱中没有找到相关信息。"
    
    # 格式化结果
    formatted_results = []
    for result in results:
        if isinstance(result, dict):
            formatted_result = f"{result.get('source', '未知源')} --{result.get('relationship', '未知关系')}--> {result.get('target', '未知目标')}"
            if result.get('source_description'):
                formatted_result += f"\n  源描述: {result['source_description']}"
            if result.get('target_description'):
                formatted_result += f"\n  目标描述: {result['target_description']}"
            formatted_results.append(formatted_result)
        else:
            formatted_results.append(str(result))
    
    return "\n\n".join(formatted_results) if formatted_results else "知识图谱中没有找到相关信息。"

@tool
def store_to_knowledge_graph(text: str) -> str:
    """将文本中的知识提取并存储到知识图谱中"""
    kg_service = get_knowledge_graph_service()
    
    # 提取实体和关系
    entities, relationships = kg_service.extract_entities_relationships(text)
    
    if not entities:
        return "未能从文本中提取到有效的实体信息。"
    
    # 存储到知识图谱
    kg_service.store_knowledge(entities, relationships)
    
    return f"成功存储 {len(entities)} 个实体和 {len(relationships)} 个关系到知识图谱。"

def get_knowledge_graph_tool():
    """获取知识图谱工具实例"""
    return KnowledgeGraphQueryTool()